#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 接收 /force_data 话题
# 如果接收到的力不超过期望力020，则控制机械臂往初始位置运动

import time
import numpy as np
import rospy
import sys
from moveit_commander import MoveGroupCommander
from geometry_msgs.msg import Pose, TwistStamped

class ForceControl:
    def __init__(self):
        self.force_torque_filtered = np.zeros((3, 1))
        self.force_torque_reference = np.zeros((3, 1))
        self.force_torque_reference[1] = 0  # 参考力
        self.position_reference = np.zeros((3, 1))
        self.tcp_frame_cur_p = np.zeros((3, 1))
        self.tcp_frame_cur_M = np.eye(3)
        self.velocity = np.zeros((3, 1))
        self.acceleration = np.zeros((3, 1))
        self.K_position = np.eye(3) * 1.1
        self.B_position = np.eye(3) * 50
        self.M_position = np.eye(3) * 280
        self.dt = 1
        self.first_get_coord = True
        self.group = MoveGroupCommander("manipulator")

        # 订阅 /force_data 话题
        self.wrench_sub = rospy.Subscriber('/force_data', TwistStamped, self.force_callback)
        
    def force_callback(self, msg):
        current_force = msg.twist.linear.z
        if (self.force_torque_filtered[0] > 0 and current_force < 0) or (self.force_torque_filtered[0] < 0 and current_force > 0):
            self.first_get_coord = True # 当正负变化时重置坐标获取标志
        self.force_torque_filtered[0] = msg.twist.linear.z*5
        self.force_torque_filtered[1] = 0
        self.force_torque_filtered[2] = 0

    def update_current_pose(self):
        current_pose = self.group.get_current_pose().pose
        self.tcp_frame_cur_p[0, 0] = current_pose.position.x
        self.tcp_frame_cur_p[1, 0] = current_pose.position.y
        self.tcp_frame_cur_p[2, 0] = current_pose.position.z
        if self.first_get_coord:
            self.position_reference[0, 0] = current_pose.position.x
            self.position_reference[1, 0] = current_pose.position.y
            self.position_reference[2, 0] = current_pose.position.z
            self.first_get_coord = False

    def calculate_new_pose(self):
        # print("z:", self.force_torque_filtered)
        force_difference = self.force_torque_filtered - self.force_torque_reference
        position_difference = self.tcp_frame_cur_p - self.position_reference
        Kp_position_term = np.dot(self.K_position, position_difference)
        B_velocity_term = np.dot(self.B_position, self.velocity)
        
        self.acceleration = np.dot(np.linalg.inv(self.M_position), 
                                   (force_difference - Kp_position_term - B_velocity_term))
        self.velocity += (self.acceleration * self.dt)
        delta_position = self.velocity * self.dt
        self.velocity = np.zeros((3, 1))
        self.velocity[0, 0] = 0.0005
        
        new_position = self.tcp_frame_cur_p - delta_position
        # 打印加速度、速度和新位置
        print("力：:\n", force_difference)
        # print("当前位置:\n", self.tcp_frame_cur_p)
        # print("参考位置:\n", self.position_reference)
        # print("位置差:\n", position_difference)
        # print("Kp_position_term:\n", Kp_position_term)
        # print("B_velocity_term:\n", B_velocity_term)
        # print("加速度:\n", self.acceleration)
        # print("Velocity:\n", self.velocity)
        # print("改变的位置量:\n", delta_position)
        # print("当前位置:\n", self.tcp_frame_cur_p)
        print("新位置:\n", new_position)
        print("\n")
        return new_position
    
    def force_below_threshold(self, threshold, tolerance=0.1):
        return np.all(np.abs(self.force_torque_filtered) <= threshold + tolerance)

allow_force_control = True  # 允许力控操作

def stop_force_control():
    global allow_force_control
    allow_force_control = False

def run_force_control():
    global allow_force_control
    allow_force_control = True

def force_main():
    # rospy.init_node('force_control_node')  # 初始化ROS节点
    group_name = "manipulator"
    move_group = MoveGroupCommander(group_name)
    force_control = ForceControl()
    print("力控开始")
    
    move_group.set_named_target('middle')
    move_group.go()
    # move_group.set_pose_target(middle)
    # move_group.go(wait=True)
    move_group.stop()
    move_group.clear_pose_targets()
    
    try:
        global allow_force_control
        while not rospy.is_shutdown() and allow_force_control:
            force_control.update_current_pose()
            # 检查接收到的力是否超过期望力
            if  force_control.force_torque_filtered[0] == 0:
                # 如果接收到的力不超过期望力，控制机械臂运动到初始位置
                move_group.set_named_target('middle')
                move_group.go(wait=False)
            else:
                new_position = force_control.calculate_new_pose()
                
                pose_target = move_group.get_current_pose().pose
                pose_target.position.x = new_position[0, 0]
                pose_target.position.y = 0
                pose_target.position.z = new_position[2, 0]

                # # 检查 z 值是否超过 0.32
                # if pose_target.position.z > 0.35 or pose_target.position.y > 0.001:
                #     if pose_target.position.z > 0.35:
                #         print("Z value exceeds 0.35. Exiting the process.")
                #     if pose_target.position.y > 0.001:
                #         print("Y value exceeds 0.001. Exiting the process.")
                #     rospy.sleep(2)
                #     move_group.set_named_target('middle')
                #     move_group.go(wait=True)
                #     break  # 结束循环和进程
                
                move_group.set_pose_target(pose_target)
                move_group.go(wait=False)
                if not allow_force_control:
                    raise KeyboardInterrupt
            time.sleep(1)
     
    except KeyboardInterrupt:
        print("Motion stopped by user")
    finally:
        move_group.stop()
        move_group.clear_pose_targets()

# if __name__ == "__main__":
#     force_main()